""" Module interp 

Exposes several functions for fitting and interpolation, 
some implemented in Cython.

Copyright 2010 (C) Almar Klein, University of Twente.

"""

# Compile cython
from pyzolib import pyximport
pyximport.install()
import interpolation_

# Import cython part
from interpolation_ import interp, project, ainterp, aproject
from interpolation_ import make_samples_absolute, fix_samples_edges
from interpolation_ import get_cubic_spline_coefs, meshgrid

# Normal imports
import numpy as np
from pirt import Point, Pointset, Aarray, is_Aarray, diffuse


## Deformations


def deform_backward(data, deltas, order=1, spline_type=0.0):
    """ deform_backward(data, deltas, order=1, spline_type=0.0)
    
    Interpolate data according to the deformations specified in deltas.
    Deltas should be a tuple of numpy arrays similar to 'samples' in
    the interp() function. They represent the relative sample positions 
    expressed in world coordinates.
    
    """
    
    # Check
    if len(deltas) != data.ndim:
        tmp = "Samples must contain as many arrays as data has dimensions."
        raise ValueError(tmp)
    
    # Create samples
    samples = make_samples_absolute(deltas)
    
    # Interpolate
    result = interp(data, samples, order, spline_type)
    
    # Make Aarray
    if is_Aarray(data):
        result = Aarray(result, data.sampling, data.origin)
    
    # Done
    return result


def deform_forward(data, deltas):
    
    # Check
    if len(deltas) != data.ndim:
        tmp = "Samples must contain as many arrays as data has dimensions."
        raise ValueError(tmp)
    
    # Create samples
    samples = make_samples_absolute(deltas)
    
    # Interpolate
    result = project(data, samples)
    
    # Make Aarray
    if is_Aarray(data):
        result = Aarray(result, data.sampling, data.origin)
    
    # Done
    return result


## Resize and zoom

def resize(data, new_shape, order=3, spline_type=0.0, prefilter=False, extra=False):
    """ resize(data, new_shape, order=3, spline_type=0.0, prefilter=False, extra=False)
    
    Resize the data to the specified new shape.
    
    Parameters
    ----------
    data : numpy array 
        The data to rezize.
    new_shape : tuple 
        The new shape of the data (z-y-x order).
    order : {0,1,3} or {'nearest', 'linear', 'cubic'}
        The interpolation order to use.
    spline_type : float or string
        Only for cubic interpolation. Specifies the type of spline. 
        Can be 'Basis', 'Hermite', 'Cardinal', 'Catmull-rom', 'Lagrange', 
        'Lanczos', 'quadratic', or a float, specifying the tension parameter
        for the Cardinal spline. See the docs of get_cubic_spline_coefs()
        for more information.
    prefilter : bool
       Whether to apply (discrete Gaussian diffusion) anti-aliasing 
       (when downampling). Default False.
    extra : bool
        Whether to extrapolate the data a bit. In this case, each datapoint
        is seen as spanning a space equal to the distance between the data
        points. This is the method used when you resize an image using 
        e.g. paint.net or photoshop. If False, the first and last datapoint
        are exactly on top of the original first and last datapoint (like
        scipy.ndimage.zoom). Default False.
    
    Notes on extrapolating
    ----------------------
    For the sake of simplicity, assume that the new shape is exactly
    twice that of the original.
    When extra if False, the sampling between the pixels is not a factor 2 
    of the original. When extra is True, the sampling decreases with a 
    factor of 2, but the data now has an offset. Additionally, extrapolation
    is performed, which is less accurate than interpolation
    
    """
    
    # Check new_shape
    if not isinstance(new_shape, (tuple,list)):
        raise ValueError('new_shape must be a tuple or list.')
    elif not len(new_shape) == len(data.shape):
        raise ValueError('new_shape must contain as much values as data has dimensions.')
    new_shape = [int(round(n)) for n in new_shape]
    
    # Get shape, sampling and origin
    shape = data.shape
    if is_Aarray(data):
        sampling = data.sampling
        origin = data.origin
    else:
        sampling = [1.0 for s in shape]
        origin = [0.0 for s in shape]
    
    # Init lists
    ranges = []
    sampling2 = []
    origin2 = []
    
    for s, n, sam, ori in zip(shape, new_shape, sampling, origin):        
        
        if extra:
            
            # Get full range (expressed in "pixels")
            dmin, dmax = -0.5, s-0.5
            drange = dmax-dmin # == s
            
            # Step size (n-1 steps in between the pixels, and 0.5 steps at each side)
            dstep = float(drange) / n  
            dstep2 = 0.5 * dstep
            
            # Get sampling and offset
            sampling2.append( dstep*sam )
            origin2.append( ori+(dmin+dstep2)*sam )
            
            # Get range
            r = np.linspace(dmin+dstep2, dmax-dstep2, n)
            ranges.append(r)
        
        else:
            
            # Get full range (expressed in "pixels")
            dmin, dmax = 0, s-1
            drange = dmax-dmin # == s
            
            # Step size (the outer points are exactly at the dmin and dmax)
            dstep = float(drange) / (n-1)
            
            # Get sampling and offset   
            sampling2.append( dstep*sam )
            origin2.append( ori )
            
            # Get range
            r = np.linspace(dmin, dmax, n)
            ranges.append(r)
    
    
    # Anti-aliasing
    def foo(x):
        if x < 1.0: return 0.8/x # 1.0 is better, but people like sharp images
        else: return 0.0
    factors = [float(s1)/s2 for s1,s2 in zip(new_shape, shape)]
    sigmas = [foo(f) for f in factors]
    if prefilter and sum(sigmas):
        data = diffuse(data, sigmas)
    
    # Interpolate (first make ranges x-y-z)
    ranges.reverse()
    grids = meshgrid(ranges)
    data2 = interp(data, grids, order, spline_type)
    
    # Make Aarray
    return Aarray(data2, sampling2, origin2)


def imresize(data, factor, order=3):
    """ imzoom(data, factor, order=3)
    
    Convenience function to resize the image data (1D, 2D or 3D).
    
    This function uses pirt.resize() with 'prefilter' and 'extra' set to True.
    This makes it more suitble for generic image resizing. Use pirt.resize()
    for more fine-grained control.
    
    Parameters
    ----------
    data : numpy array 
        The data to rezize.
    new_shape : tuple 
        The new shape of the data (z-y-x order).
    order : {0,1,3} or {'nearest', 'linear', 'cubic'}
        The interpolation order to use.
    
    """
    return imresize(data, new_shape, order, 0.0, True, True)


def zoom(data, factor, order=3, spline_type=0.0, prefilter=False, extra=False):
    """ zoom(data, factor, order=3, spline_type=0.0, prefilter=False, extra=False)
    
    Resize the data with the specified factor. The default behavior is
    the same as scipy.ndimage.zoom(), but three times faster.
    
    Parameters
    ----------
    data : numpy array 
        The data to rezize.
    factor : scalar or tuple 
        The resize factor, optionally for each dimension (z-y-z order).
    order : {0,1,3} or {'nearest', 'linear', 'cubic'}
        The interpolation order to use.
    spline_type : float or string
        Only for cubic interpolation. Specifies the type of spline. 
        Can be 'Basis', 'Hermite', 'Cardinal', 'Catmull-rom', 'Lagrange', 
        'Lanczos', 'quadratic', or a float, specifying the tension parameter
        for the Cardinal spline. See the docs of get_cubic_spline_coefs()
        for more information.
    prefilter : bool
       Whether to apply (discrete Gaussian diffusion) anti-aliasing 
       (when downampling). Default False.
    extra : bool
        Whether to extrapolate the data a bit. In this case, each datapoint
        is seen as spanning a space equal to the distance between the data
        points. This is the method used when you resize an image using 
        e.g. paint.net or photoshop. If False, the first and last datapoint
        are exactly on top of the original first and last datapoint (like
        numpy.zoom). Default False.
    
    Notes on extrapolating
    ----------------------
    For the sake of simplicity, assume a resize factor of 2.
    When extra if False, the sampling between the pixels is not a factor 2
    of the original. When extra is True, the sampling decreases with a
    factor of 2, but the data now has an offset. Additionally, extrapolation
    is performed, which is less accurate than interpolation
    
    """
   
    # Process factor
    if isinstance(factor, np.ndarray) and np.size == 1:
        factor = float(factor)
    if isinstance(factor, (float, int)):
        factor = [factor for i in data.shape]
    
    # Check factor
    if not isinstance(factor, (list, tuple)):
        raise ValueError('Factor must be a float or tuple/list.')
    if len(factor) != data.ndim:
        raise ValueError('Factor len does not match ndim of data.')
    
    # Calculate new shape
    new_shape = [float(f)*s for f,s in zip(factor, data.shape)]
    new_shape = [int(round(s)) for s in new_shape]
    
    # Resize
    return resize(data, new_shape, order, spline_type, prefilter, extra)


def imzoom(data, factor, order=3):
    """ imzoom(data, factor, order=3)
    
    Convenience function to resize the image data (1D, 2D or 3D) with the
    specified factor.
    
    This function uses pirt.resize() with 'prefilter' and 'extra' set to True.
    This makes it more suitble for generic image resizing. Use pirt.resize()
    for more fine-grained control.
    
    Parameters
    ----------
    data : numpy array 
        The data to rezize.
    factor : scalar or tuple 
        The resize factor, optionally for each dimension (z-y-x order).
    order : {0,1,3} or {'nearest', 'linear', 'cubic'}
        The interpolation order to use.
    
    """
    return zoom(data, factor, order, 0.0, True, True)


## For slices in 3D volumes
def get_span_vectors(normal, c, d):
    """ get_span_vectors(normal, prevA, prevB) -> (a,b)
    
    Given a normal, return two orthogonal vectors which are both orthogonal
    to the normal. The vectors are calculated so they match as much as possible
    the previous vectors.
    
    """
    
    # Calculate a from previous b
    a1 = d.cross(normal)
    
    if a1.norm() < 0.001:
        # The normal and  d point in same or reverse direction
        # -> Calculate b from previous a
        b1 = c.cross(normal)
        a1 = b1.cross(normal)
    
    # Consider the opposite direction
    a2 = -1 * a1
    if c.distance(a1) > c.distance(a2):
        a1 = a2
    
    # Ok, calculate b
    b1 = a1.cross(normal)
    
    # Consider the opposite
    b2 = -1 * b1
    if d.distance(b1) > d.distance(b2):
        b1 = b2

    # Done
    return a1.normalize(), b1.normalize()


class SliceInVolume:
    """ SliceInVolume(self, pos, normal=None, previous=None)
    Defines a slice in a volume. 
    
    The two span vectors are in v and u respectively. In other words,
    vec1 is up, vec2 is right.
    """
    
    def __init__(self, pos, normal=None, previous=None):
        
        # Init vectors
        self._pos = pos
        self._normal = None
        self._vec1 = None
        self._vec2 = None
        
        # Calculate normal
        if normal is None and previous is None:
            self._normal = Point(0,0,-1)
        elif normal is None:
            # Normal is defined by difference in position,
            # and the previous normal
            self._normal = pos - previous._pos            
            self._normal += previous._normal
            self._normal = self._normal.normalize()
        elif normal is not None:
            self._normal = normal.normalize()
        
        # Calculate vec1 and vec2
        if previous is None:
            # Use arbitrary vector
            arbitrary = Point(1,0,0)
            self._vec1 = arbitrary.cross(self._normal)
            if self._vec1.norm() < 0.0001:
                arbitrary = Point(0,1,0)
                self._vec1 = arbitrary.cross(self._normal)
            # second
            self._vec2 = self._vec1.cross(self._normal)
        else:
            # Use previous
            tmp = get_span_vectors(self._normal, previous._vec1, previous._vec2)
            self._vec1, self._vec2 = tmp
        # Normalize
        self._vec1 = self._vec1.normalize()
        self._vec2 = self._vec2.normalize()
    
    def get_slice(self, volume, N=128, spacing=1.0):
        vec1 = self._vec1 * spacing
        vec2 = self._vec2 * spacing
        im = interpolation_.slice_from_volume(volume, self._pos, vec1, vec2, N)
        return Aarray(im, (spacing,spacing))
    
    
    def convert_local_to_global(self, p2d, p3d):
        """ convert_local_to_global(p2d, p3d)
        Convert local 2D points to global 3D points.
        UNTESTED
        """
        pos = self._pos.copy()
        #
        pos.x += p2d.y * self._vec1.x + p2d.x * self._vec2.x
        pos.y += p2d.y * self._vec1.y + p2d.x * self._vec2.y
        pos.z += p2d.y * self._vec1.z + p2d.x * self._vec2.z
        #
        return pos
